import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import FuncAnimation
from satellite_orbit import CircularOrbit

# 定义模拟参数
simulation_time = 1000 * 60  # 模拟 1000 分钟，转换为秒
time_step = 60  # 时间步长为 1 分钟，转换为秒
num_steps = int(simulation_time / time_step)

# 创建一个 3D 图形
fig = plt.figure(figsize=(10, 10))
ax = fig.add_subplot(111, projection='3d')

# 设置地球参数
EARTH_RADIUS = 6371.0
u = np.linspace(0, 2 * np.pi, 100)
v = np.linspace(0, np.pi, 100)
x_earth = EARTH_RADIUS * np.outer(np.cos(u), np.sin(v))
y_earth = EARTH_RADIUS * np.outer(np.sin(u), np.sin(v))
z_earth = EARTH_RADIUS * np.outer(np.ones(np.size(u)), np.cos(v))

# 绘制地球（灰色透明球体）
ax.plot_surface(x_earth, y_earth, z_earth, color='gray', alpha=0.2)

# 定义多个卫星轨道
orbits = [
    CircularOrbit(orbit_radius=8000.0, inclination=np.pi / 4, raan=0.0, satellite_count=8, color='r'),
    CircularOrbit(orbit_radius=9000.0, inclination=np.pi / 3, raan=np.pi / 2, satellite_count=8, color='g'),
    CircularOrbit(orbit_radius=10000.0, inclination=np.pi / 6, raan=np.pi / 4, satellite_count=8, color='b')
]

# 绘制每个轨道
orbit_lines = []
satellite_points = []
normal_vectors = []
for i, orbit in enumerate(orbits):
    orbit_points = orbit.get_orbit_points()
    x_orbit = [point[0] for point in orbit_points]
    y_orbit = [point[1] for point in orbit_points]
    z_orbit = [point[2] for point in orbit_points]
    line, = ax.plot(x_orbit, y_orbit, z_orbit, color=orbit.color, label=f'Orbit {i + 1}')
    orbit_lines.append(line)

    satellites = ax.scatter([], [], [], color=orbit.color, s=50)
    satellite_points.append(satellites)

    # 计算并绘制法向量
    normal = orbit.get_normal_vector()
    origin = [0, 0, 0]
    normal_line = ax.quiver(*origin, *normal, length=orbit.orbit_radius, color=orbit.color, label=f'Normal {i + 1}')
    normal_vectors.append(normal_line)

# 设置标题和图例
ax.set_title('Satellite Orbits around the Earth')
ax.legend()

# 设置坐标轴标签
ax.set_xlabel('X (km)')
ax.set_ylabel('Y (km)')
ax.set_zlabel('Z (km)')

# 设置坐标轴比例为 1:1:1
ax.set_box_aspect([1, 1, 1])


# 更新函数，用于动画的每一帧
def update(frame):
    time = frame * time_step
    for i, orbit in enumerate(orbits):
        satellite_positions = orbit.get_all_satellite_positions(time)
        x_sat = [pos[0] for pos in satellite_positions]
        y_sat = [pos[1] for pos in satellite_positions]
        z_sat = [pos[2] for pos in satellite_positions]
        satellite_points[i]._offsets3d = (x_sat, y_sat, z_sat)

        # 更新法向量（由于法向量不随时间变化，这里可以省略更新操作）
        # normal = orbit.get_normal_vector()
        # origin = [0, 0, 0]
        # normal_vectors[i].set_segments([[origin, [origin[j] + normal[j] * orbit.orbit_radius for j in range(3)]]])

    return satellite_points + normal_vectors


# 创建动画
ani = FuncAnimation(fig, update, frames=num_steps, interval=200, blit=False)

# 显示动画
plt.show()